package com.lahuan.config.datasouce;

import com.lahuan.config.datasouce.entity.GlobalXid;
import com.lahuan.config.datasouce.entity.XaBranch;
import com.mysql.cj.jdbc.MysqlXADataSource;
import com.mysql.cj.jdbc.MysqlXid;
import com.mysql.cj.jdbc.SuspendableXAConnection;

import javax.sql.DataSource;
import javax.transaction.xa.XAResource;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.util.UUID;

/**
 * @author lfy
 */
public class XaProxyUtil {
    /**
     * 代理数据源，在获取链接的时候，转换成xa连接
     * @param ds
     * @return
     */
    public static DataSource proxyDatasource(MysqlXADataSource ds ,String username,String password,String dsName){

        return (DataSource) Proxy.newProxyInstance(Object.class.getClassLoader(),
                new Class[]{DataSource.class}, (proxy, method, args) -> {
                    String name = method.getName();
                    System.out.println(dsName+":"+name);
                    if("getConnection".equals(name)){
                        GlobalXid globalXid = XaTransactionHolder.getGlobalXid();
                        if(globalXid ==null){
                            //没有XA事务时
                            return method.invoke(ds, args);
                        }
                        //有XA事务时转换成xa连接
                        SuspendableXAConnection xaConnection = (SuspendableXAConnection) ds.getXAConnection(username,password);

                        XAResource rm = xaConnection.getXAResource();
                        int formatId = globalXid.getFormatId();
                        byte[] gtrid = globalXid.getGtrid();
                        byte[] bqual = UUID.randomUUID().toString().getBytes();
                        MysqlXid mysqlXid = new MysqlXid(gtrid, bqual, formatId);
                        System.out.println(dsName+":start xa");
                        rm.start(mysqlXid, XAResource.TMNOFLAGS);
                        XaBranch branch = new XaBranch();
                        branch.setMysqlXid(mysqlXid);
                        branch.setRm(rm);
                        XaTransactionHolder.addXaBranchSet(branch);
                        XaTransactionHolder.setCurrentXaBranchHolder(branch);
                        return proxyConnection(xaConnection.getConnection(),dsName);
                    }
                    return method.invoke(ds, args);
                });
    }

    public static Connection proxyConnection(Object connection,String dsName){
        return (Connection) Proxy.newProxyInstance(Object.class.getClassLoader(),
                new Class[]{Connection.class}, (proxy, method, args) -> {
                    String name = method.getName();
                    System.out.println(dsName+"-connection:"+name);
                    if("close".equals(name)){
                        method.invoke(connection, args);
                        XaBranch branch = XaTransactionHolder.getCurrentXaBranch();
                        MysqlXid mysqlXid = branch.getMysqlXid();
                        XAResource rm = branch.getRm();
                        System.out.println(dsName+":end xa");
                        rm.end(mysqlXid, XAResource.TMSUCCESS);
                        return null;
                    }
                    return method.invoke(connection, args);
                });
    }

}
